# coding: utf-8
import os

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
import torch
import pandas as pd
from torchvision import transforms
from PIL import Image
from config_hyperparam import cfg
from model import swin_base_patch4_window12_384_in22k as Model
from torchvision.transforms import InterpolationMode
from torch import nn
import json
import warnings
warnings.filterwarnings("ignore")

def load_class_labels(json_file):
    with open(json_file, 'r') as f:
        class_labels = json.load(f)
    index_to_class = {v: k for k, v in class_labels.items()}
    return index_to_class


def load_model(model_path):
    model = Model(cfg.num_class).to(cfg.device)
    checkpoint = torch.load(model_path)
    model.load_state_dict(checkpoint["model"])
    model.eval()  
    return model


def predict_image(model, image_path, index_to_class, transform):
    image = Image.open(image_path).convert('RGB')
    if transform:
        image = transform(image)
    image = image.unsqueeze(0)
    image = image.to(cfg.device)
    with torch.no_grad(): 
        outputs, _ = model(image)
        _, predicted = torch.max(nn.Softmax(dim=1)(outputs), 1) 
        
    predicted_index = predicted.item()
    with open('translations.json', 'r', encoding='utf-8') as f:
        trans = json.load(f)
    return trans.get(index_to_class.get(predicted_index))

if __name__ == "__main__":

    model_path = cfg.best_model_path  
    csv_file = 'class_labels.json' 
    image_dir = './demo'  


    index_to_class = load_class_labels(csv_file)


    model = load_model(model_path)


    transform = transforms.Compose([
        transforms.Resize((256, 256), InterpolationMode.BILINEAR),
        transforms.CenterCrop((cfg.resize, cfg.resize)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    print("----------------- 开始识别样本 -----------------")
    for image_name in os.listdir(image_dir):
        if image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            image_path = os.path.join(image_dir, image_name)
            predicted_class = predict_image(model, image_path, index_to_class, transform)
            print("样本:{}, 识别品种:{}".format(image_name, predicted_class))
            print("-------------------------------------")
    print("----------------- 识别结束 -----------------")